﻿// Copyright (c) 2025 liucl
// Inject is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

#include "enum_wnd.hpp"

#include <Windows.h>
#include <tlhelp32.h>

#include <boost/scope_exit.hpp>

#include <vector>

std::vector<DWORD> get_thread_ids(DWORD pid) {
  std::vector<DWORD> thread_ids;
  // 创建线程快照
  HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, pid);
  if (hSnapshot == INVALID_HANDLE_VALUE) {
    return thread_ids; // 返回空向量
  }
  BOOST_SCOPE_EXIT(hSnapshot) { CloseHandle(hSnapshot); }
  BOOST_SCOPE_EXIT_END;

  THREADENTRY32 te;
  memset(&te, 0, sizeof(THREADENTRY32));
  te.dwSize = sizeof(THREADENTRY32);
  for (auto ret = Thread32First(hSnapshot, &te); ret;
       ret = Thread32Next(hSnapshot, &te)) {
    // 检查线程是否属于指定进程
    if (te.th32OwnerProcessID == pid) {
      thread_ids.push_back(te.th32ThreadID);
    }
  }
  return thread_ids;
}

bool enum_windows(WNDS &wnds) {
  auto pid = ::GetCurrentProcessId();
  // 获取进程所有线程
  auto thread_ids = get_thread_ids(pid);
  // 枚举每个线程的窗口
  for (auto thread_id : thread_ids) {
    ::EnumThreadWindows(
        thread_id,
        [](HWND hwnd, LPARAM lParam) -> BOOL {
          auto &wnds = *reinterpret_cast<WNDS *>(lParam);
          wnds.push_back(hwnd);
          return true;
        },
        reinterpret_cast<LPARAM>(&wnds));
  }

  return !wnds.empty();
}
